import argparse
import os
import lpips
import clip_score

parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('--dir0', type=str, default='D:\\PycharmProjects\\IDM-main\\experiments\\coco_scaler\\results\\hr')
parser.add_argument('--dir1', type=str, default='D:\\PycharmProjects\\IDM-main\\experiments\\coco_scaler\\results\\sr')
parser.add_argument('-v', '--version', type=str, default='0.1')
opt = parser.parse_args()

## Initializing the model
loss_fn = lpips.LPIPS(net='alex', version=opt.version)

# the total list of images
files = os.listdir(opt.dir0)
files_2 = os.listdir(opt.dir1)
i = 0
total_lpips_distance = 0
average_lpips_distance = 0
for rname, fname in zip(files, files_2):

    try:
        # Load images
        img0 = lpips.im2tensor(lpips.load_image(os.path.join(opt.dir0, rname)))
        img1 = lpips.im2tensor(lpips.load_image(os.path.join(opt.dir1, fname)))

        if (os.path.exists(os.path.join(opt.dir0, rname)), os.path.exists(os.path.join(opt.dir1, fname))):
            i = i + 1

        # Compute distance
        current_lpips_distance = loss_fn.forward(img0, img1)
        total_lpips_distance = total_lpips_distance + current_lpips_distance

        print('%s: %.3f' % (rname, current_lpips_distance))

    except Exception as e:
        print(e)

average_lpips_distance = float(total_lpips_distance) / i

print("The processed iamges is ", i, "and the average_lpips_distance is: %.3f" % average_lpips_distance)

# python -m clip_score C:\\Users\\admin\\Desktop\\ours\\hr D:\\PycharmProjects\\Image-Super-Resolution-via-Iterative-Refinement-master\\dataset\\test_500_caption --device cuda:0 --num-workers 0
